from cryptography.fernet import Fernet
from django.db import models
from django.core.files.base import ContentFile
import os


def get_encryption_key():
    return EncryptedFileField.fernet_key


class EncryptedFileField(models.FileField):
    # fernet_key = Fernet.generate_key()
    fernet_key = b'GyPq_tY3W9ldTG5Ln7cNrKI9EPf6GjQisssCk0qXyM0='
    cipher_suite = Fernet(fernet_key)

    def __init__(self, *args, **kwargs):
        # self.fernet_key = fernet_key
        # self.cipher_suite = cipher_suite
        super().__init__(*args, **kwargs)

    def pre_save(self, model_instance, add):
        file = getattr(model_instance, self.attname)

        # 只在文件对象存在且未提交时处理
        if file and not file._committed:
            # 读取文件内容
            content = file.read()

            # 加密文件内容
            encrypted_content = self.cipher_suite.encrypt(content)

            # 创建加密文件的内容对象
            encrypted_file = ContentFile(encrypted_content)

            # 保存加密文件
            # file_name = self.generate_filename(model_instance, os.path.basename(file.name))
            if hasattr(model_instance, 'image'):
                file_name = model_instance.image.name
            elif hasattr(model_instance, 'video'):
                file_name = model_instance.video.name
            # temp = model_instance.media_root
            # file.path = os.path.join(temp,file_name)
            file.save(file_name, encrypted_file, save=False)

        return super().pre_save(model_instance, add)
    
    def get_encryption_key():
        return EncryptedFileField.fernet_key


class EncryptFile:
    # fernet_key = Fernet.generate_key()
    fernet_key = b'GyPq_tY3W9ldTG5Ln7cNrKI9EPf6GjQisssCk0qXyM0='
    cipher_suite = Fernet(fernet_key)

    def encrypt(self, file):
        with open(file, 'rb') as f:
            content = f.read()
            # 加密文件内容
            encrypted_content = self.cipher_suite.encrypt(content)
            # 创建加密文件的内容对象

        with open(file, 'wb') as f:
            f.write(encrypted_content)



    def get_encryption_key(self):
        return EncryptedFileField.fernet_key
